#! /usr/bin/env python3
import logging
import platform
import subprocess
import sys
import os
import stat
import shutil
from argparse import ArgumentParser, RawDescriptionHelpFormatter
import argparse
from collections import OrderedDict
import atexit
import json
import shlex

# Packages which we wish to ensure are always recompiled
wheel_blacklist = ["mpi4py", "randomgen", "islpy"]

# Firedrake application installation shortcuts.
firedrake_apps = {
    "gusto": ("""Atmospheric dynamical core library. http://firedrakeproject.org/gusto""",
              "git+ssh://github.com/firedrakeproject/gusto#egg=gusto"),
    "thetis": ("""Coastal ocean model. http://thetisproject.org""",
               "git+ssh://github.com/thetisproject/thetis#egg=thetis"),
    "icepack": ("""Glacier and ice sheet model. https://icepack.github.io""",
                "git+ssh://github.com/icepack/icepack.git#egg=icepack"),
}


class InstallError(Exception):
    # Exception for generic install problems.
    pass


class FiredrakeConfiguration(dict):
    """A dictionary extended to facilitate the storage of Firedrake
    configuration information."""
    def __init__(self, args=None):
        super(FiredrakeConfiguration, self).__init__()

        '''A record of the persistent options in force.'''
        self["options"] = {}
        '''Relevant environment variables.'''
        self["environment"] = {}
        '''Additional packages installed via the plugin interface.'''
        self["additions"] = []

        if args:
            for o in self._persistent_options:
                if o in args.__dict__.keys():
                    self["options"][o] = args.__dict__[o]

    _persistent_options = ["package_manager",
                           "minimal_petsc", "mpicc", "mpicxx", "mpif90", "mpiexec", "disable_ssh",
                           "honour_petsc_dir", "with_parmetis",
                           "slepc", "packages", "honour_pythonpath",
                           "opencascade",
                           "petsc_int_type", "cache_dir"]


def deep_update(this, that):
    from collections import abc
    for k, v in that.items():
        if isinstance(v, abc.Mapping) and k in this.keys():
            this[k] = deep_update(this.get(k, {}), v)
        else:
            this[k] = v
    return this


if os.path.basename(__file__) == "firedrake-install":
    mode = "install"
    logfile_directory = os.path.abspath(os.getcwd())
    logfile_mode = "w"

elif os.path.basename(__file__) == "firedrake-update":
    mode = "update"
    os.chdir(os.path.dirname(os.path.realpath(__file__)) + "/../..")
    try:
        logfile_directory = os.environ["VIRTUAL_ENV"]
    except KeyError:
        quit("Unable to retrieve venv name from the environment.\n Please ensure the venv is active before running firedrake-update.")
    logfile_mode = "a" if "--no-update-script" in sys.argv else "w"
else:
    sys.exit("Script must be invoked either as firedrake-install or firedrake-update")


# Set up logging
# Log to file at DEBUG level
logfile = os.path.join(logfile_directory, 'firedrake-%s.log' % mode)
logging.basicConfig(level=logging.DEBUG,
                    format='%(asctime)s %(levelname)-6s %(message)s',
                    filename=logfile,
                    filemode=logfile_mode)
# Log to console at INFO level
console = logging.StreamHandler()
console.setLevel(logging.INFO)
formatter = logging.Formatter('%(message)s')
console.setFormatter(formatter)
logging.getLogger().addHandler(console)

log = logging.getLogger()

log.info("Running %s" % " ".join(sys.argv))


if sys.version_info < (3, 5):
    if mode == "install":
        print("""\nInstalling Firedrake requires Python 3, at least version 3.5.
You should run firedrake-install with python3.""")
    if mode == "update":
        if hasattr(sys, "real_prefix"):
            # sys.real_prefix exists iff we are in an active virtualenv.
            #
            # Existing install trying to update past the py2/py3 barrier
            print("""\nFiredrake is now Python 3 only.  You cannot upgrade your existing installation.
Please follow the instructions at http://www.firedrakeproject.org/download.html to reinstall.""")
            sys.exit(1)
        else:
            # Accidentally (?) running firedrake-update with python2.
            print("""\nfiredrake-update must be run with Python 3, did you accidentally use Python 2?""")
    sys.exit(1)


branches = {}

jenkins = "JENKINS_URL" in os.environ
ci_testing_firedrake = "FIREDRAKE_CI_TESTS" in os.environ


def resolve_doi_branches(doi):
    import requests
    import hashlib
    log.info("Installing Firedrake components specified by {}".format(doi))
    response = requests.get("https://zenodo.org/api/records",
                            params={"q": "doi:{}".format(doi.replace("/", r"\/"))})
    if response.status_code >= 400:
        log.error("Unable to obtain Zenodo record for doi {}".format(doi))
        log.error("Response was {}".format(response.json()))
        sys.exit(1)
    response = response.json()
    try:
        record, = response["hits"]["hits"]
    except ValueError:
        log.error("Was expecting one record for doi '{doi}', found {num}".format(
            doi=doi, num=response["hits"]["total"]))
        log.error("Response was {}".format(response))
        sys.exit(1)
    files = record["files"]
    try:
        componentjson, = (f for f in files if f["key"] == "components.json")
    except ValueError:
        log.error("Expecting to find exactly one 'components.json' in record")
        sys.exit(1)
    download = requests.get(componentjson["links"]["self"])
    if download.status_code >= 400:
        log.error("Unable to download 'components.json'")
        log.error("Response was {}".format(download.json()))
        sys.exit(1)
    # component response has checksum as "md5:HEXDIGEST", strip the md5.
    if hashlib.md5(download.content).hexdigest() != componentjson["checksum"][4:]:
        log.error("Download failed checksum, expecting {expect}, got {got}".format(
            expect=componentjson["checksum"][4:],
            got=hashlib.md5(download.content).hexdigest()))
        sys.exit(1)
    componentjson = download.json()
    branches = {}
    for record in componentjson["components"]:
        commit = record["commit"]
        component = record["component"]
        package = component[component.find("/")+1:].lower()
        branches[package] = commit
        log.info("Using commit {commit} for component {comp}".format(
            commit=commit, comp=component))
    return branches


if mode == "install":
    # Handle command line arguments.
    parser = ArgumentParser(description="""Install firedrake and its dependencies.""",
                            epilog="""The install process has three steps.

1. Any required system packages are installed using brew (MacOS) or apt (Ubuntu
   and similar Linux systems). On a Linux system without apt, the installation
   will fail if a dependency is not found.

2. A set of standard and/or third party Python packages is installed to the
   specified install location.

3. The core set of Python packages is downloaded to ./firedrake/src/ and
   installed to the specified location.

The install creates a venv in ./firedrake and installs inside
that venv.

The installer will ensure that the required configuration options are
passed to PETSc. In addition, any configure options which you provide
in the PETSC_CONFIGURE_OPTIONS environment variable will be
honoured.""",
                            formatter_class=RawDescriptionHelpFormatter)

    parser.add_argument("--slepc", action="store_true",
                        help="Install SLEPc along with PETSc.")
    parser.add_argument("--opencascade", action="store_true",
                        help="Install OpenCASCADE for CAD integration.")
    parser.add_argument("--disable-ssh", action="store_true",
                        help="Do not attempt to use ssh to clone git repositories: fall immediately back to https.")
    parser.add_argument("--no-package-manager", action='store_false', dest="package_manager",
                        help="Do not attempt to use apt or homebrew to install operating system packages on which we depend.")
    group = parser.add_mutually_exclusive_group()
    group.add_argument("--minimal-petsc", action="store_true",
                       help="Minimise the set of petsc dependencies installed. This creates faster build times (useful for build testing).")
    group.add_argument("--with-parmetis", action="store_true",
                       help="Install PETSc with ParMETIS? (Non-free license, see http://glaros.dtc.umn.edu/gkhome/metis/parmetis/download)")
    group = parser.add_mutually_exclusive_group()
    group.add_argument("--honour-petsc-dir", action="store_true",
                       help="Usually it is best to let Firedrake build its own PETSc. If you wish to use another PETSc, set PETSC_DIR and pass this option.")
    group.add_argument("--petsc-int-type", choices=["int32", "int64"],
                       default="int32", type=str,
                       help="The integer type used by PETSc.  Use int64 if you need to solve problems with more than 2 billion degrees of freedom.  Only takes effect if firedrake-install builds PETSc.")

    parser.add_argument("--honour-pythonpath", action="store_true",
                        help="Pointing to external Python packages is usually a user error. Set this option if you know that you want PYTHONPATH set.")
    parser.add_argument("--rebuild-script", action="store_true",
                        help="Only rebuild the firedrake-install script. Use this option if your firedrake-install script is broken and a fix has been released in upstream Firedrake. You will need to specify any other options which you wish to be honoured by your new update script.")
    parser.add_argument("--doi", type=str, nargs=1,
                        help="Install a set of components matching a particular Zenodo DOI. The record should have been created with firedrake-zenodo.")
    # Used for testing if Zenodo broke the API
    # Tries to resolve to a known DOI and immediately exits.
    parser.add_argument("--test-doi-resolution", action="store_true",
                        help=argparse.SUPPRESS)
    parser.add_argument("--package-branch", type=str, nargs=2, action="append", metavar=("PACKAGE", "BRANCH"),
                        help="Specify which branch of a package to use. This takes two arguments, the package name and the branch.")
    parser.add_argument("--verbose", "-v", action="store_true", help="Produce more verbose debugging output.")
    parser.add_argument("--mpicc", type=str,
                        action="store", default=None,
                        help="C compiler to use when building with MPI. If not set, MPICH will be downloaded and used.")
    parser.add_argument("--mpicxx", type=str,
                        action="store", default=None,
                        help="C++ compiler to use when building with MPI. If not set, MPICH will be downloaded and used.")
    parser.add_argument("--mpif90", type=str,
                        action="store", default=None,
                        help="Fortran compiler to use when building with MPI. If not set, MPICH will be downloaded and used.")
    parser.add_argument("--mpiexec", type=str,
                        action="store", default=None,
                        help="MPI launcher. If not set, MPICH will be downloaded and used.")
    parser.add_argument("--show-petsc-configure-options", action="store_true",
                        help="Print out the configure options passed to PETSc and exit")
    parser.add_argument("--show-dependencies", action="store_true",
                        help="Print out the package manager used and packages installed/required.")
    parser.add_argument("--venv-name", default="firedrake",
                        type=str, action="store",
                        help="Name of the venv to create (default is 'firedrake')")
    parser.add_argument("--install", action="append", dest="packages",
                        help="Additional packages to be installed. The address should be in format vcs+protocol://repo_url/#egg=pkg&subdirectory=pkg_dir . Some additional packages have shortcut install options for more information see --install-help.")
    parser.add_argument("--install-help", action="store_true",
                        help="Provide information on packages which can be installed using shortcut names.")
    parser.add_argument("--cache-dir", type=str,
                        action="store",
                        help="Directory to use for disk caches of compiled code (default is the .cache subdirectory of the Firedrake installation).")
    parser.add_argument("--documentation-dependencies", action="store_true",
                        help="Install the dependencies required to build the documentation")

    args = parser.parse_args()

    # If the user has set any MPI info, they must set them all
    if args.mpicc or args.mpicxx or args.mpif90 or args.mpiexec:
        if not (args.mpicc and args.mpicxx and args.mpif90 and args.mpiexec):
            log.error("If you set any MPI information, you must set all of {mpicc, mpicxx, mpif90, mpiexec}.")
            sys.exit(1)

    if args.package_branch:
        branches = {package.lower(): branch for package, branch in args.package_branch}
    if args.test_doi_resolution:
        actual = resolve_doi_branches("10.5281/zenodo.1322546")
        expect = {'coffee': '87e50785d3a05b111f5423a66d461cd44cc4bdc9',
                  'finat': 'aa74fd499304c8363a4520555fd62ef21e8e5e1f',
                  'fiat': '184601a46c24fb5cbf8fd7961d22b16dd26890e7',
                  'firedrake': '6a30b64da01eb587dcc0e04e8e6b84fe4839bdb7',
                  'petsc': '413f72f04f5cb0a010c85e03ed029573ff6d4c63',
                  'petsc4py': 'ac2690070a80211dfdbab04634bdb3496c14ca0a',
                  'tsfc': 'fe9973eaacaa205fd491cd1cc9b3743b93a3d076',
                  'ufl': 'c5eb7fbe89c1091132479c081e6fa9c182191dcc',
                  'pyop2': '741a21ba9a62cb67c0aa300a2e199436ea8cb61c'}
        if actual != expect:
            log.error("Unable to resolve DOI correctly.")
            log.error("You'll need to figure out how Zenodo have changed their API.")
            sys.exit(1)
        else:
            log.info("DOI resolution test passed.")
            sys.exit(0)
    if args.doi:
        branches = resolve_doi_branches(args.doi[0])

    args.prefix = False  # Disabled as untested
    args.packages = args.packages or []

    config = FiredrakeConfiguration(args)
else:
    # This duplicates code from firedrake_configuration in order to avoid the module dependency and allow for installation recovery.
    try:
        with open(os.path.join(os.environ["VIRTUAL_ENV"],
                               ".configuration.json"), "r") as f:
            config = json.load(f)

    except FileNotFoundError:
        # Fall back to the old location.
        import firedrake_configuration

        config = firedrake_configuration.get_config()
        if config is None:
            raise InstallError("Failed to find existing Firedrake configuration")

    parser = ArgumentParser(description="""Update this firedrake install to the latest versions of all packages.""",
                            formatter_class=RawDescriptionHelpFormatter)
    parser.add_argument("--no-update-script", action="store_false", dest="update_script",
                        help="Do not update script before updating Firedrake.")
    parser.add_argument("--rebuild", action="store_true",
                        help="Rebuild all packages even if no new version is available. Usually petsc and petsc4py are only rebuilt if they change. All other packages are always rebuilt.")
    parser.add_argument("--rebuild-script", action="store_true",
                        help="Only rebuild the firedrake-install script. Use this option if your firedrake-install script is broken and a fix has been released in upstream Firedrake. You will need to specify any other options which you wish to be honoured by your new update script.")
    parser.add_argument("--slepc", action="store_true", dest="slepc", default=config["options"]["slepc"],
                        help="Install SLEPc along with PETSc")
    parser.add_argument("--opencascade", action="store_true", dest="opencascade", default=config["options"].get("opencascade", False),
                        help="Install OpenCASCADE for CAD integration.")
    group = parser.add_mutually_exclusive_group()
    group.add_argument("--honour-petsc-dir", action="store_true",
                       default=config["options"].get("honour_petsc_dir", False),
                       help="Usually it is best to let Firedrake build its own PETSc. If you wish to use another PETSc, set PETSC_DIR and pass this option.")
    group.add_argument("--petsc-int-type", choices=["int32", "int64"],
                       default="int32", type=str,
                       help="The integer type used by PETSc.  Use int64 if you need to solve problems with more than 2 billion degrees of freedom.  Only takes effect if firedrake-install builds PETSc.")

    parser.add_argument("--honour-pythonpath", action="store_true", default=config["options"].get("honour_pythonpath", False),
                        help="Pointing to external Python packages is usually a user error. Set this option if you know that you want PYTHONPATH set.")
    parser.add_argument("--clean", action='store_true',
                        help="Delete any remnants of obsolete Firedrake components.")
    parser.add_argument("--verbose", "-v", action="store_true", help="Produce more verbose debugging output.")
    parser.add_argument("--install", action="append", dest="packages", default=config["options"].get("packages", []),
                        help="Additional packages to be installed. The address should be in format vcs+protocol://repo_url/#egg=pkg&subdirectory=pkg_dir. Some additional packages have shortcut install options for more information see --install-help.")
    parser.add_argument("--install-help", action="store_true",
                        help="Provide information on packages which can be installed using shortcut names.")
    parser.add_argument("--cache-dir", type=str,
                        action="store", default=config["options"].get("cache_dir", ""),
                        help="Directory to use for disk caches of compiled code (default is the .cache subdirectory of the Firedrake installation).")
    parser.add_argument("--documentation-dependencies", action="store_true",
                        help="Install the dependencies required to build the documentation")
    parser.add_argument("--show-dependencies", action="store_true",
                        help="Print out the package manager used and packages installed/required.")

    args = parser.parse_args()

    args.packages = list(set(args.packages))  # remove duplicates

    petsc_int_type_changed = False
    if config["options"].get("petsc_int_type", "int32") != args.petsc_int_type:
        petsc_int_type_changed = True
        args.rebuild = True

    config = deep_update(config, FiredrakeConfiguration(args))

if args.install_help:
    help_string = """
You can install the following packages by passing --install shortname
where shortname is one of the names given below:

"""
    componentformat = "|{:10}|{:70}|\n"
    header = componentformat.format("Name", "Description")
    line = "-" * (len(header) - 1) + "\n"
    help_string += line + header + line
    for package, d in firedrake_apps.items():
        help_string += componentformat.format(package, d[0])
    help_string += line
    print(help_string)
    sys.exit(0)


@atexit.register
def print_log_location():
    log.info("\n\n%s log saved in %s" % (mode.capitalize(), logfile))


class directory(object):
    """Context manager that executes body in a given directory"""
    def __init__(self, dir):
        self.dir = os.path.abspath(dir)

    def __enter__(self):
        self.olddir = os.path.abspath(os.getcwd())
        log.debug("Old path '%s'" % self.olddir)
        log.debug("Pushing path '%s'" % self.dir)
        os.chdir(self.dir)

    def __exit__(self, *args):
        log.debug("Popping path '%s'" % self.dir)
        os.chdir(self.olddir)
        log.debug("New path '%s'" % self.olddir)


class environment(object):
    def __init__(self, **env):
        self.old = os.environ.copy()
        self.new = env

    def __enter__(self):
        os.environ.update(self.new)

    def __exit__(self, *args):
        os.environ = self.old


options = config["options"]

# Apply short cut package names. pyadjoint shoud no longer turn up the packages list because it is a hard dependency.
options["packages"] = [firedrake_apps.get(p, (None, p))[1] for p in options["packages"] if not p.endswith("pyadjoint")]


# Record of obsolete packages which --clean should remove from old installs.
old_git_packages = ["dolfin-adjoint", "libadjoint"]

if mode == "install":
    firedrake_env = os.path.abspath(args.venv_name)
else:
    firedrake_env = os.environ["VIRTUAL_ENV"]

if "cache_dir" not in config["options"] or not config["options"]["cache_dir"]:
    config["options"]["cache_dir"] = os.path.join(firedrake_env, ".cache")

# venv install
python = ["%s/bin/python" % firedrake_env]
# Use the pip from the venv
pip = python + ["-m", "pip"]
pipinstall = pip + ["install", "--no-binary", ",".join(wheel_blacklist)]


# This context manager should be used whenever arguments should be temporarily added to pipinstall.
class pipargs(object):
    def __init__(self, *args):
        self.args = args

    def __enter__(self):
        self.save = pipinstall.copy()
        pipinstall.extend(self.args)

    def __exit__(self, *args, **kwargs):
        global pipinstall
        pipinstall = self.save


pyinstall = python + ["setup.py", "install"]

if "PYTHONPATH" in os.environ and not args.honour_pythonpath:
    quit("""The PYTHONPATH environment variable is set. This is probably an error.
If you really want to use your own Python packages, please run again with the
--honour-pythonpath option.
""")

petsc_opts = "--download-eigen=%s/src/eigen-3.3.3.tgz " % firedrake_env

petsc_opts += "--with-fortran-bindings=0 "

petsc_options = {"--with-fortran-bindings=0",
                 "--with-debugging=0",
                 "--with-shared-libraries=1",
                 "--with-c2html=0",
                 "--download-eigen=%s/src/eigen-3.3.3.tgz " % firedrake_env,
                 # File format
                 "--download-hdf5",
                 # AMG
                 "--download-hypre",
                 # Sparse direct solver
                 "--download-superlu_dist",
                 # Parallel mesh partitioner
                 "--download-ptscotch",
                 # For superlu_dist amongst others.
                 "--with-cxx-dialect=C++11"}

if not options["minimal_petsc"]:
    petsc_options.add("--with-zlib")
    # File formats
    petsc_options.add("--download-netcdf")
    petsc_options.add("--download-pnetcdf")
    petsc_options.add("--download-exodusii")
    # Sparse direct solvers
    petsc_options.add("--download-suitesparse")
    petsc_options.add("--download-pastix")
    # Required by pastix
    petsc_options.add("--download-hwloc")
    # Serial mesh partitioner
    petsc_options.add("--download-metis")
    if options.get("with_parmetis", None):
        # Non-free license.
        petsc_options.add("--download-parmetis")

if options["petsc_int_type"] == "int32":
    # Sparse direct solver
    petsc_options.add("--download-scalapack")  # Needed for mumps
    petsc_options.add("--download-mumps")
    # Serial mesh partitioner
    petsc_options.add("--download-chaco")
    if not options["minimal_petsc"]:
        # AMG
        petsc_options.add("--download-ml")
else:
    petsc_options.add("--with-64-bit-indices")

if options.get("mpiexec") is not None:
    petsc_options.add("--with-mpiexec={}".format(options["mpiexec"]))
    petsc_options.add("--with-cc={}".format(options["mpicc"]))
    petsc_options.add("--with-cxx={}".format(options["mpicxx"]))
    petsc_options.add("--with-fc={}".format(options["mpif90"]))
else:
    # Download mpich if the user does not tell us about an MPI.
    petsc_options.add("--download-mpich")

petsc_options = list(petsc_options) + shlex.split(os.environ.get("PETSC_CONFIGURE_OPTIONS", ""))

if mode == "update" and petsc_int_type_changed:
    log.warning("""Force rebuilding all packages because PETSc int type changed""")


def check_call(arguments):
    try:
        log.debug("Running command '%s'", " ".join(arguments))
        log.debug(subprocess.check_output(arguments, stderr=subprocess.STDOUT, env=os.environ).decode())
    except subprocess.CalledProcessError as e:
        log.debug(e.output.decode())
        raise


def check_output(args):
    try:
        log.debug("Running command '%s'", " ".join(args))
        return subprocess.check_output(args, stderr=subprocess.STDOUT, env=os.environ).decode()
    except subprocess.CalledProcessError as e:
        log.debug(e.output.decode())
        raise


def brew_install(name, options=None):
    try:
        # Check if it's already installed
        check_call(["brew", "list", name])
    except subprocess.CalledProcessError:
        # If not found, go ahead and install
        arguments = [name]
        if options:
            arguments = options + arguments
        if args.verbose:
            arguments = ["--verbose"] + arguments
        check_call(["brew", "install"] + arguments)


def apt_check(name):
    log.info("Checking for presence of package %s..." % name)
    # Note that subprocess return codes have the opposite logical
    # meanings to those of Python variables.
    try:
        check_call(["dpkg-query", "-s", name])
        log.info("  installed.")
        return True
    except subprocess.CalledProcessError:
        log.info("  missing.")
        return False


def apt_install(names):
    log.info("Installing missing packages: %s." % ", ".join(names))
    if sys.stdin.isatty():
        subprocess.check_call(["sudo", "apt-get", "install"] + names)
    else:
        log.info("Non-interactive stdin detected; installing without prompts")
        subprocess.check_call(["sudo", "apt-get", "-y", "install"] + names)


def split_requirements_url(url):
    name = url.split(".git")[0].split("#")[0].split("/")[-1]
    spliturl = url.split("://")[1].split("#")[0].split("@")
    try:
        plain_url, branch = spliturl
    except ValueError:
        plain_url = spliturl[0]
        branch = "master"
    return name, plain_url, branch


def git_url(plain_url, protocol):
    if protocol == "ssh":
        return "git@%s:%s" % tuple(plain_url.split("/", 1))
    elif protocol == "https":
        return "https://%s" % plain_url
    else:
        raise ValueError("Unknown git protocol: %s" % protocol)


def git_clone(url):
    name, plain_url, branch = split_requirements_url(url)
    if name == "petsc" and args.honour_petsc_dir:
        log.info("Using existing petsc installation\n")
        return name
    log.info("Cloning %s\n" % name)
    branch = branches.get(name.lower(), branch)
    try:
        if options["disable_ssh"]:
            raise InstallError("Skipping ssh clone because --disable-ssh")
        # note: so far only loopy requires submodule
        check_call(["git", "clone", "-q", "--recursive", git_url(plain_url, "ssh")])
        log.info("Successfully cloned repository %s" % name)
    except (subprocess.CalledProcessError, InstallError):
        if not options["disable_ssh"]:
            log.warning("Failed to clone %s using ssh, falling back to https." % name)
        try:
            check_call(["git", "clone", "-q", "--recursive", git_url(plain_url, "https")])
            log.info("Successfully cloned repository %s." % name)
        except subprocess.CalledProcessError:
            log.error("Failed to clone %s branch %s." % (name, branch))
            raise
    with directory(name):
        try:
            log.info("Checking out branch %s" % branch)
            check_call(["git", "checkout", "-q", branch])
            log.info("Successfully checked out branch %s" % branch)
        except subprocess.CalledProcessError:
            log.error("Failed to check out branch %s" % branch)
            raise
    return name


def list_cloned_dependencies(name):
    log.info("Finding dependencies of %s\n" % name)
    deps = OrderedDict()
    try:
        for dep in open(name + "/requirements-git.txt", "r"):
            name = split_requirements_url(dep.strip())[0]
            deps[name] = dep.strip()
    except IOError:
        pass
    return deps


def clone_dependencies(name):
    log.info("Cloning the dependencies of %s" % name)
    deps = []
    try:
        for dep in open(name + "/requirements-git.txt", "r"):
            deps.append(git_clone(dep.strip()))
    except IOError:
        pass
    return deps


def git_update(name, url=None):
    # Update the named git repo and return true if the current branch actually changed.
    log.info("Updating the git repository for %s" % name)
    with directory(name):
        git_sha = check_output(["git", "rev-parse", "HEAD"])
        # Ensure remotes get updated if and when we move repositories.
        if url:
            plain_url = split_requirements_url(url)[1]
            current_url = check_output(["git", "remote", "-v"]).split()[1]
            protocol = "https" if current_url.startswith("https") else "ssh"
            new_url = git_url(plain_url, protocol)
            # Ensure we only change from bitbucket to github and not the reverse.
            if new_url != current_url and "bitbucket.org" in current_url \
               and ("github.com/firedrakeproject" in plain_url
                    or "github.com/dolfin-adjoint" in plain_url):
                log.info("Updating git remote for %s" % name)
                check_call(["git", "remote", "set-url", "origin", new_url])
        check_call(["git", "pull", "--recurse-submodules"])
        git_sha_new = check_output(["git", "rev-parse", "HEAD"])
    return git_sha != git_sha_new


def run_pip(args):
    check_call(pip + args)


def run_pip_install(pipargs):
    # Make pip verbose when logging, so we see what the
    # subprocesses wrote out.
    # Particularly important for debugging petsc fails.
    pipargs = ["-vvv"] + pipargs
    check_call(pipinstall + pipargs)


def run_cmd(args):
    check_call(args)


def pip_requirements(package):
    log.info("Installing pip dependencies for %s" % package)
    if os.path.isfile("%s/requirements-ext.txt" % package):
        run_pip_install(["-r", "%s/requirements-ext.txt" % package])
    elif os.path.isfile("%s/requirements.txt" % package):
        if package == "COFFEE":
            # FIXME: Horrible hack to work around
            # https://github.com/coin-or/pulp/issues/123
            run_pip_install(["--no-deps", "-r", "%s/requirements.txt" % package])
        elif package == "pyadjoint":
            with open("%s/requirements.txt" % package, "r") as f:
                reqs = f.readlines()
            for req in reqs:
                req = req.strip()
                try:
                    run_pip_install([req])
                except subprocess.CalledProcessError as e:
                    if req == "tensorflow":
                        log.debug("Skipping failing install of optional pyadjoint dependency tensorflow")
                    else:
                        raise e
        elif package == "loopy":
            # loopy dependencies are installed in setup.py
            return
        else:
            run_pip_install(["-r", "%s/requirements.txt" % package])
    else:
        log.info("No dependencies found. Skipping.")


def install(package):
    log.info("Installing %s" % package)
    if package == "petsc/":
        build_and_install_petsc()
    # The following outrageous hack works around the fact that petsc and co. cannot be installed in developer mode.
    elif package not in ["petsc4py/", "slepc/", "slepc4py/"]:
        run_pip_install(["-e", package])
    else:
        run_pip_install(["--ignore-installed", package])


def clean(package):
    log.info("Cleaning %s" % package)
    with directory(package):
        check_call(["python", "setup.py", "clean"])


def pip_uninstall(package):
    log.info("Removing existing %s installations\n" % package)
    # Uninstalling something with pip is an absolute disaster.  We
    # have to use pip freeze to list all available packages "locally"
    # and keep on removing the one we want until it is gone from this
    # list!  Yes, pip will happily have two different versions of the
    # same package co-existing.  Moreover, depending on the phase of
    # the moon, the order in which they are uninstalled is not the
    # same as the order in which they appear on sys.path!
    again = True
    uninstalled = False
    i = 0
    while again:
        # List installed packages, "locally".  In a venv,
        # this just tells me packages in the venv, otherwise it
        # gives me everything.
        lines = check_output(pip + ["freeze", "-l"])
        again = False
        for line in lines.split("\n"):
            # Do we have a locally installed package?
            if line.startswith(package+'=='):
                # Uninstall it.
                run_pip(["uninstall", "-y", line.strip()])
                uninstalled = True
                # Go round again, because THERE MIGHT BE ANOTHER ONE!
                again = True
        i += 1
        if i > 10:
            raise InstallError("pip claims it uninstalled %s more than 10 times.  Something is probably broken.", package)
    return uninstalled


def get_petsc_dir():
    if args.honour_petsc_dir:
        try:
            petsc_dir = os.environ["PETSC_DIR"]
        except KeyError:
            raise InstallError("Unable to find installed PETSc (did you forget to set PETSC_DIR?)")
        petsc_arch = os.environ.get("PETSC_ARCH", "")
    else:
        petsc_dir = os.path.join(os.environ["VIRTUAL_ENV"], "src", "petsc")
        petsc_arch = "default"

    return petsc_dir, petsc_arch


def get_slepc_dir():
    petsc_dir, petsc_arch = get_petsc_dir()
    if args.honour_petsc_dir:
        try:
            slepc_dir = os.environ["SLEPC_DIR"]
        except KeyError:
            raise InstallError("Need to set SLEPC_DIR for --slepc with --honour-petsc-dir")
    else:
        try:
            slepc_dir = check_output(python
                                     + ["-c", "import slepc; print(slepc.get_slepc_dir())"]).strip()
        except subprocess.CalledProcessError:
            raise InstallError("Unable to find installed SLEPc")
    return slepc_dir, petsc_arch


def build_and_install_petsc():
    import hashlib
    import urllib.request
    tarball = "eigen-3.3.3.tgz"
    url = "https://github.com/eigenteam/eigen-git-mirror/archive/3.3.3.tar.gz"
    sha = hashlib.sha256()
    # Accept eigen either from github or the previous bitbucket location.
    expect = ("df6082711532336354466bab70f66c16ff88e616ac3884eba555b6f0b4860e65",
              "94878cbfa27b0d0fbc64c00d4aafa137f678d5315ae62ba4aecddbd4269ae75f")
    if not os.path.exists(tarball):
        log.info("Downloading Eigen from '%s' to '%s'" % (url, tarball))
        urllib.request.urlretrieve(url, filename=tarball)
    else:
        log.info("Eigen tarball already downloaded to '%s'" % tarball)
    log.info("Checking Eigen tarball integrity")
    with open(tarball, "rb") as f:
        while True:
            data = f.read(65536)
            if not data:
                break
            sha.update(data)
    actual = sha.hexdigest()
    if actual not in expect:
        raise InstallError("Downloaded Eigen tarball has incorrect sha256sum, expected '%s', was '%s'",
                           expect, actual)
    else:
        log.info("Eigen tarball hash valid")
    log.info("Depending on your platform, PETSc may take an hour or more to build!")
    assert not args.honour_petsc_dir  # Belt and braces security before we start clobbering things.
    with directory("petsc"):
        petsc_dir, petsc_arch = get_petsc_dir()
        check_call(["rm", "-rf", os.path.join(petsc_dir, petsc_arch)])
        check_call(["./configure", "PETSC_DIR={}".format(petsc_dir), "PETSC_ARCH={}".format(petsc_arch)] + list(petsc_options))
        check_call(["make", "PETSC_DIR={}".format(petsc_dir), "PETSC_ARCH={}".format(petsc_arch), "all"])


def build_and_install_h5py():
    import shutil
    log.info("Installing h5py")
    # Clean up old downloads
    if os.path.exists("h5py-2.5.0"):
        log.info("Removing old h5py-2.5.0 source")
        shutil.rmtree("h5py-2.5.0")
    if os.path.exists("h5py.tar.gz"):
        log.info("Removing old h5py.tar.gz")
        os.remove("h5py.tar.gz")

    url = "git+https://github.com/firedrakeproject/h5py.git@firedrake#egg=h5py"
    if os.path.exists("h5py"):
        changed = False
        with directory("h5py"):
            # Rewrite old h5py/h5py remote to firedrakeproject/h5py remote.
            plain_url = split_requirements_url(url)[1]
            current_remote = check_output(["git", "remote", "-v"]).split()[1]
            proto = "https" if current_remote.startswith("https") else "ssh"
            new_remote = git_url(plain_url, proto)
            if new_remote != current_remote and "h5py/h5py.git" in current_remote:
                log.info("Updating git remote for h5py from %s to %s", current_remote, new_remote)
                check_call(["git", "remote", "set-url", "origin", new_remote])
                check_call(["git", "fetch", "-p", "origin"])
                check_call(["git", "checkout", "firedrake"])
                changed = True
        changed |= git_update("h5py")
    else:
        git_clone(url)
        changed = True

    petsc_dir, petsc_arch = get_petsc_dir()
    hdf5_dir = "%s/%s" % (petsc_dir, petsc_arch)
    if changed or args.rebuild:
        log.info("Linking h5py against PETSc found in %s\n" % hdf5_dir)
        with environment(HDF5_DIR=hdf5_dir,
                         HDF5_MPI="ON"):
            # Only uninstall if things changed.
            pip_uninstall("h5py")
            # Pip installing from dirty directory is potentially unsafe.
            with directory("h5py/"):
                check_call(["git", "clean", "-fdx"])
            install("h5py/")
    else:
        log.info("No need to rebuild h5py")


def build_and_install_libsupermesh():
    log.info("Installing libsupermesh")
    url = "git+https://bitbucket.org/libsupermesh/libsupermesh.git"
    if os.path.exists("libsupermesh"):
        log.info("Updating the git repository for libsupermesh")
        with directory("libsupermesh"):
            check_call(["git", "fetch"])
            git_sha = check_output(["git", "rev-parse", "HEAD"])
            git_sha_new = check_output(["git", "rev-parse", "@{u}"])
            changed = git_sha != git_sha_new
            if changed:
                check_call(["git", "reset", "--hard"])
                check_call(["git", "pull"])
    else:
        git_clone(url)
        changed = True
    if changed:
        with directory("libsupermesh"):
            check_call(["git", "reset", "--hard"])
            check_call(["git", "clean", "-f", "-x", "-d"])
            check_call(["mkdir", "-p", "build"])
            with directory("build"):
                check_call(["cmake", "..", "-DBUILD_SHARED_LIBS=ON",
                            "-DCMAKE_INSTALL_PREFIX=" + firedrake_env,
                            "-DMPI_C_COMPILER=" + cc,
                            "-DMPI_CXX_COMPILER=" + cxx,
                            "-DMPI_Fortran_COMPILER=" + f90,
                            "-DCMAKE_Fortran_COMPILER=" + f90])
                check_call(["make"])
                check_call(["make", "install"])
    else:
        log.info("No need to rebuild libsupermesh")


def build_and_install_pythonocc():
    log.info("Installing pythonocc-core")
    url = "git+https://github.com/tpaviot/pythonocc-core.git@595b0a4e8e60e8d6011bea0cdb54ac878efcfcd2"
    if not os.path.exists("pythonocc-core"):
        git_clone(url)
        with directory("pythonocc-core"):
            check_call(["git", "reset", "--hard"])
            check_call(["git", "clean", "-f", "-x", "-d"])
            check_call(["mkdir", "-p", "build"])
            with directory("build"):
                check_call(["cmake", "..",
                            "-DCMAKE_INSTALL_PREFIX=" + firedrake_env,
                            "-DPYTHON_EXECUTABLE=" + python[0],
                            "-DMPI_C_COMPILER=" + cc,
                            "-DMPI_CXX_COMPILER=" + cxx,
                            "-DMPI_Fortran_COMPILER=" + f90,
                            "-DCMAKE_Fortran_COMPILER=" + f90])
                check_call(["make", "install"])
    else:
        log.info("No need to rebuild pythonocc-core")


def build_and_install_libspatialindex():
    log.info("Installing libspatialindex")
    if os.path.exists("libspatialindex"):
        log.info("Updating the git repository for libspatialindex")
        with directory("libspatialindex"):
            check_call(["git", "fetch"])
            git_sha = check_output(["git", "rev-parse", "HEAD"])
            git_sha_new = check_output(["git", "rev-parse", "@{u}"])
            libspatialindex_changed = git_sha != git_sha_new
            if libspatialindex_changed:
                check_call(["git", "reset", "--hard"])
                check_call(["git", "pull"])
    else:
        git_clone("git+https://github.com/firedrakeproject/libspatialindex.git")
        libspatialindex_changed = True

    if libspatialindex_changed:
        with directory("libspatialindex"):
            # Clean source directory
            check_call(["git", "reset", "--hard"])
            check_call(["git", "clean", "-f", "-x", "-d"])
            # Patch Makefile.am to skip building test
            check_call(["sed", "-i", "-e", "/^SUBDIRS/s/ test//", "Makefile.am"])
            # Build and install
            check_call(["./autogen.sh"])
            check_call(["./configure", "--prefix=" + firedrake_env,
                        "--enable-shared", "--disable-static"])
            check_call(["make"])
            check_call(["make", "install"])
    else:
        log.info("No need to rebuild libspatialindex")


def build_and_install_slepc():
    petsc_dir, petsc_arch = get_petsc_dir()

    if args.honour_petsc_dir:
        slepc_dir, slepc_arch = get_slepc_dir()
        log.info("Using installed SLEPc from %s/%s", slepc_dir, slepc_arch)
    else:
        log.info("Installing SLEPc.")
        url = "git+https://github.com/firedrakeproject/slepc.git@firedrake"
        if os.path.exists("slepc"):
            slepc_changed = False
            with directory("slepc"):
                plain_url = split_requirements_url(url)[1]
                current_remote = check_output(["git", "remote", "-v"]).split()[1]
                proto = "https" if current_remote.startswith("https") else "ssh"
                new_remote = git_url(plain_url, proto)
                if new_remote != current_remote:
                    log.info("Updating git remote for SLEPc from %s to %s", current_remote, new_remote)
                    check_call(["git", "remote", "set-url", "origin", new_remote])
                    check_call(["git", "fetch", "-p", "origin"])
                    check_call(["git", "checkout", "firedrake"])
                    slepc_changed = True
            slepc_changed |= git_update("slepc")
        else:
            git_clone(url)
            slepc_changed = True
        if slepc_changed:
            install("slepc/")
        else:
            log.info("No need to rebuild SLEPc")

    log.info("Installing slepc4py.")
    url = "git+https://github.com/firedrakeproject/slepc4py.git@firedrake"
    if os.path.exists("slepc4py"):
        slepc4py_changed = False
        with directory("slepc4py"):
            plain_url = split_requirements_url(url)[1]
            current_remote = check_output(["git", "remote", "-v"]).split()[1]
            proto = "https" if current_remote.startswith("https") else "ssh"
            new_remote = git_url(plain_url, proto)
            if new_remote != current_remote:
                log.info("Updating git remote for slepc4py from %s to %s", current_remote, new_remote)
                check_call(["git", "remote", "set-url", "origin", new_remote])
                check_call(["git", "fetch", "-p", "origin"])
                check_call(["git", "checkout", "firedrake"])
                slepc4py_changed = True
        slepc4py_changed |= git_update("slepc4py")
    else:
        git_clone(url)
        slepc4py_changed = True
    if slepc4py_changed:
        install("slepc4py/")
    else:
        log.info("No need to rebuild slepc4py")


def install_documentation_dependencies():
    """Just install the required dependencies. There are no provenance
    issues here so no need to record this in the configuration dict."""
    log.info("Installing documentation dependencies")
    run_pip_install(["sphinx"])
    run_pip_install(["sphinxcontrib-bibtex"])
    run_pip_install(["git+https://github.com/sphinx-contrib/youtube.git"])


def clean_obsolete_packages():
    dirs = os.listdir(".")
    for package in old_git_packages:
        pip_uninstall(package)
        if package in dirs:
            shutil.rmtree(package)


def quit(message):
    log.error(message)
    sys.exit(1)


def build_update_script():
    log.info("Creating firedrake-update script.")
    with open("firedrake/scripts/firedrake-install", "r") as f:
        update_script = f.read()

    try:
        os.mkdir("../bin")
    except OSError:
        pass
    with open("../bin/firedrake-update", "w") as f:
        f.write(update_script)
    check_call(["chmod", "a+x", "../bin/firedrake-update"])


if args.rebuild_script:
    os.chdir(os.path.dirname(os.path.realpath(__file__)) + ("/../.."))

    build_update_script()

    log.info("Successfully rebuilt firedrake-update.\n")

    log.info("To upgrade firedrake, run firedrake-update")
    sys.exit(0)


def create_compiler_env(cc, cxx, f90):
    env = dict()
    if cc:
        env["MPICC"] = cc
        env["MPI_C_COMPILER"] = cc
        env["CC"] = cc
        # Work around homebrew adding /usr/local/lib to the library search path.
        env["CFLAGS"] = " ".join((os.environ.get("CFLAGS", ""),
                                  "-L" + os.path.join(*get_petsc_dir(), "lib"),
                                  "-I" + os.path.join(*get_petsc_dir(), "include")))
    if cxx:
        env["MPICXX"] = cxx
        env["MPI_CXX_COMPILER"] = cxx
        env["CXX"] = cxx
    if f90:
        env["MPIF90"] = f90
        env["MPI_C_COMPILER"] = f90
        env["F90"] = f90

    return env


if mode == "install" and args.show_petsc_configure_options:
    log.info("*********************************************")
    log.info("Would build PETSc with the following options:")
    log.info("*********************************************\n")
    log.info("\n".join(petsc_options))
    log.info("\nEigen will be downloaded from https://github.com/eigenteam/eigen-git-mirror/archive/3.3.3.tar.gz")
    sys.exit(0)


if "PETSC_DIR" in os.environ and not args.honour_petsc_dir:
    quit("""The PETSC_DIR environment variable is set. This is probably an error.
If you really want to use your own PETSc build, please run again with the
--honour-petsc-dir option.
""")

if "PETSC_DIR" not in os.environ and args.honour_petsc_dir:
    quit("""The --honour-petsc-dir is set, but PETSC_DIR environment variable is
not defined. If you have compiled PETSc manually, set PETSC_DIR
(and optionally PETSC_ARCH) variables to point to the build directory.
""")

if "SLEPC_DIR" not in os.environ and args.honour_petsc_dir and options["slepc"]:
    quit("""If you use --honour-petsc-dir, you must also build SLEPc manually
and set the SLEPC_DIR environment variable appropriately""")

if "SLEPC_DIR" in os.environ and not args.honour_petsc_dir and options["slepc"]:
    quit("""The SLEPC_DIR environment variable is set.  If you want to use
your own SLEPc version, you must also build your own PETSc and run
with --honour-petsc-dir.""")

if platform.uname()[0] == "Darwin" and args.opencascade:
    quit("""Sorry, automatically installing opencascade on OSX hasn't been
implemented yet. (It's not a supported package in brew.)
Please contact us to get this working.""")

log.debug("*** Current environment (output of 'env') ***")
log.debug(check_output(["env"]))
log.debug("\n\n")

if mode == "install" or not args.update_script:
    # Check operating system.
    osname = platform.uname()[0]
    required_packages = None
    package_manager = None
    if osname == "Darwin":
        package_manager = "brew"
        required_packages = ["gcc",
                             "python3",
                             "autoconf",
                             "automake",
                             "cmake",
                             "libtool",
                             "mercurial",
                             "boost"]
    elif osname == "Linux":
        package_manager = "apt-get"
        required_packages = ["build-essential",
                             "autoconf",
                             "automake",
                             "bison",  # for ptscotch
                             "flex",   # for ptscotch
                             "cmake",
                             "gfortran",
                             "git",
                             "libblas-dev",
                             "liblapack-dev",
                             "libtool",
                             "mercurial",
                             "mpich",
                             "python3-dev",
                             "python3-pip",
                             "python3-tk",
                             "python3-venv",
                             "zlib1g-dev",
                             "libboost-dev"]  # For ROL in pyadjoint
        if args.opencascade:
            required_packages.append("liboce-ocaf-dev")
            required_packages.append("swig")
    else:
        log.warning("You do not appear to be running Linux or MacOS. Please do not be surprised if this install fails.")

    if args.show_dependencies:
        log.info("\n************************************************")
        log.info("The following dependencies need to be installed:")
        log.info("************************************************\n")
        log.info("* A C and C++ compiler (for example gcc/g++ or clang), GNU make")
        log.info("* A Fortran compiler (for PETSc)")
        log.info("* Blas and Lapack")
        log.info("* Boost")
        log.info("* Git, Mercurial")
        log.info("* Python version >=3.5")
        log.info("* The Python headers")
        log.info("* autoconf, automake, libtool")
        log.info("* CMake")
        if osname == "Linux":
            log.info("* zlib, bison, flex")
        if package_manager:
            log.info("\n********************************************")
            log.info("The following %s packages are recommended:" % (package_manager))
            log.info("********************************************\n")
            log.info("\n".join(required_packages))
            log.info("\nThey will be installed by %s unless --no-package-manager is used." % (package_manager))
            if osname == "Darwin":
                log.info("\nIn addition, the command line tools will be installed using xcode")
        sys.exit(0)

    if osname == "Darwin":
        if options["package_manager"]:

            log.info("Installing command line tools...")
            try:
                check_call(["xcode-select", "--install"])
            except subprocess.CalledProcessError:
                # expected failure if already installed
                pass

            try:
                check_call(["brew", "--version"])
            except subprocess.CalledProcessError:
                quit("Homebrew not found. Please install it using the instructions at http://brew.sh and then try again.")

            log.info("Installing required packages via homebrew. You can safely ignore warnings that packages are already installed")
            # Ensure a fortran compiler is available
            for package in required_packages:
                brew_install(package)
        else:
            log.info("Xcode and homebrew installation disabled. Proceeding on the rash assumption that packaged dependencies are in place.")
            log.info("Please use the --show-dependencies for more information.")

    elif osname == "Linux":
        # Check for apt.
        try:
            if not options["package_manager"]:
                raise InstallError
            check_call(["apt-get", "--version"])

            missing_packages = [p for p in required_packages if not apt_check(p)]
            if missing_packages:
                apt_install(missing_packages)

        except (subprocess.CalledProcessError, InstallError):
            log.info("apt-get not found or disabled. Proceeding on the rash assumption that your compiled dependencies are in place.")
            log.info("Please use the --show-dependencies for more information.")


if mode == "install":
    if os.path.exists(firedrake_env):
        log.warning("Specified venv '%s' already exists", firedrake_env)
        quit("Can't install into existing venv '%s'" % firedrake_env)

    log.info("Creating firedrake venv in '%s'." % firedrake_env)
    # Debian's Python3 is screwed, they don't ship ensurepip as part
    # of the base python package, so the default python -m venv
    # doesn't work.  Moreover, they have spiked the file such that it
    # calls sys.exit, which will kill any attempts to create a venv
    # with pip.
    try:
        import ensurepip        # noqa: F401
        with_pip = True
    except ImportError:
        with_pip = False
    import venv
    venv.EnvBuilder(with_pip=with_pip).create(firedrake_env)
    if not with_pip:
        import urllib.request
        log.debug("ensurepip unavailable, bootstrapping pip using get-pip.py")
        urllib.request.urlretrieve("https://bootstrap.pypa.io/get-pip.py", filename="get-pip.py")
        check_call(python + ["get-pip.py"])
        log.debug("bootstrapping pip succeeded")
        log.debug("Removing get-pip.py")
        os.remove("get-pip.py")
    # Ensure pip and setuptools are at the latest version.
    run_pip(["install", "-U", "setuptools"])
    run_pip(["install", "-U", "pip"])

    # We haven't activated the venv so we need to manually set the environment.
    os.environ["VIRTUAL_ENV"] = firedrake_env

petsc_dir, petsc_arch = get_petsc_dir()

# Set up the MPI wrappers
for opt in ["mpicc", "mpicxx", "mpif90"]:
    if options.get(opt):
        name = options[opt]
        src = shutil.which(name)
    else:
        src = os.path.join(petsc_dir, petsc_arch, "bin", opt)
    dest = os.path.join(firedrake_env, "bin", opt)
    log.debug("Creating a symlink from %s to %s" % (src, dest))
    if os.path.lexists(dest):
        os.remove(dest)
    os.symlink(src, dest)

if options.get("mpiexec"):
    mpiexec_loc = shutil.which(options.get("mpiexec"))
else:
    mpiexec_loc = os.path.join(petsc_dir, petsc_arch, "bin", "mpiexec")
mpiexecf = os.path.join(firedrake_env, "bin", "mpiexec")
log.debug("Creating an mpiexec wrapper for %s" % mpiexec_loc)
if os.path.exists(mpiexecf):
    os.remove(mpiexecf)
with open(mpiexecf, "w") as mpiexec:
    contents = '#!/bin/bash' + os.linesep + mpiexec_loc + ' "$@"'
    mpiexec.write(contents)
os.chmod(mpiexecf, os.stat(mpiexecf).st_mode | stat.S_IEXEC)

cc = options["mpicc"]
cxx = options["mpicxx"]
f90 = options["mpif90"]
compiler_env = create_compiler_env(cc, cxx, f90)
os.chdir(firedrake_env)

# Configuration is now complete, and we have a venv so we dump the
# config file in order to maximise the chance that firedrake-update
# can be run on a partial install.
config_location = os.path.join(firedrake_env, ".configuration.json")
json_output = json.dumps(config)
with open(config_location, "w") as f:
    f.write(json_output)
log.info("Configuration saved to " + config_location)


if mode == "install":
    os.mkdir("src")
    os.chdir("src")

    if jenkins and ci_testing_firedrake:
        check_call(["ln", "-s", "../../../", "firedrake"])
    else:
        git_clone("git+https://github.com/firedrakeproject/firedrake.git")

    # Build the update script as soon as possible so that installation
    # recovery can be attempted if required.
    build_update_script()

    packages = clone_dependencies("firedrake")
    packages = clone_dependencies("PyOP2") + packages
    packages += ["firedrake"]

    for p in options["packages"]:
        name = git_clone(p)
        packages.extend(clone_dependencies(name))
        packages += [name]

    if args.honour_petsc_dir:
        packages.remove("petsc")

    # Force Cython to install first to work around pip dependency issues.
    run_pip_install(["Cython>=0.22"])

    # Need to install petsc first in order to resolve hdf5 dependency.
    if not args.honour_petsc_dir:
        with environment(**compiler_env):
            with pipargs("--no-deps"):
                packages.remove("petsc")
                install("petsc/")
        os.environ["PETSC_DIR"] = petsc_dir
        os.environ["PETSC_ARCH"] = petsc_arch
    os.environ["HDF5_DIR"] = os.path.join(petsc_dir, petsc_arch)
    os.environ["NETCDF4_DIR"] = os.path.join(petsc_dir, petsc_arch)

    # If petsc built mpich then we now need to set the compiler environment variables.
    if not compiler_env:
        compilerbin = os.path.join(petsc_dir, petsc_arch, "bin")
        cc = os.path.join(compilerbin, "mpicc")
        cxx = os.path.join(compilerbin, "mpicxx")
        f90 = os.path.join(compilerbin, "mpif90")
        compiler_env = create_compiler_env(cc, cxx, f90)

    with environment(**compiler_env):
        for p in packages:
            pip_requirements(p)

        build_and_install_h5py()
        build_and_install_libspatialindex()
        build_and_install_libsupermesh()

        if "loopy" in packages:
            # We do want to install loopy's dependencies.
            install("loopy/")
            packages.remove("loopy")

        with pipargs("--no-deps"):
            for p in packages:
                install(p+"/")
                sys.path.append(os.getcwd() + "/" + p)

        # Work around easy-install.pth bug.
        try:
            packages.remove("petsc")
        except ValueError:
            pass
        packages.remove("petsc4py")
        packages.remove("firedrake")

else:
    # Update mode
    os.chdir("src")

    if args.update_script:
        # Pull firedrake, rebuild update script, launch new script
        git_update("firedrake")
        build_update_script()
        os.execv(sys.executable, [sys.executable, "../bin/firedrake-update", "--no-update-script"] + sys.argv[1:])

    deps = OrderedDict()
    deps.update(list_cloned_dependencies("PyOP2"))
    deps.update(list_cloned_dependencies("firedrake"))
    for p in options["packages"]:
        name = split_requirements_url(p)[0]
        deps.update(list_cloned_dependencies(name))
        deps[name] = p
    packages = list(deps.keys())
    packages += ["firedrake"]

    # update packages.
    if not args.honour_petsc_dir:
        petsc_changed = git_update("petsc", deps["petsc"])
    else:
        petsc_changed = False
    petsc4py_changed = git_update("petsc4py", deps["petsc4py"])

    packages.remove("petsc")
    packages.remove("petsc4py")

    if args.clean:
        clean_obsolete_packages()
        for package in packages:
            pip_uninstall(package)
        if args.rebuild:
            pip_uninstall("petsc4py")

    # If there is a petsc package to remove, then we're an old installation which will need to rebuild PETSc.
    petsc_changed = pip_uninstall("petsc") or petsc_changed

    for p in packages:
        try:
            git_update(p, deps.get(p, None))
        except OSError as e:
            if e.errno == 2:
                log.warning("%s missing, cloning anew.\n" % p)
                git_clone(deps[p])
            else:
                raise

    with pipargs("--no-deps"):
        with environment(**compiler_env):
            # Only rebuild petsc if it has changed.
            if not args.honour_petsc_dir and (args.rebuild or petsc_changed):
                clean("petsc/")
                install("petsc/")
        if not args.honour_petsc_dir:
            os.environ["PETSC_DIR"] = petsc_dir
            os.environ["PETSC_ARCH"] = petsc_arch
        os.environ["HDF5_DIR"] = os.path.join(petsc_dir, petsc_arch)
        os.environ["NETCDF4_DIR"] = os.path.join(petsc_dir, petsc_arch)

        # If petsc built mpich then we now need to set the compiler environment variables.
        if not compiler_env:
            compilerbin = os.path.join(petsc_dir, petsc_arch, "bin")
            cc = os.path.join(compilerbin, "mpicc")
            cxx = os.path.join(compilerbin, "mpicxx")
            f90 = os.path.join(compilerbin, "mpif90")
            compiler_env = create_compiler_env(cc, cxx, f90)

    with environment(**compiler_env):
        # update dependencies.
        for p in packages:
            pip_requirements(p)

        with pipargs("--no-deps"):
            if args.rebuild or petsc_changed or petsc4py_changed:
                clean("petsc4py/")
                install("petsc4py/")

            # Always rebuild h5py.
            build_and_install_h5py()
            build_and_install_libspatialindex()
            build_and_install_libsupermesh()

    if "loopy" in packages:
        # We do want to install loopy's dependencies.
        clean("loopy")
        install("loopy/")
        packages.remove("loopy")
    with pipargs("--no-deps"):
        try:
            packages.remove("PyOP2")
            packages.remove("firedrake")
        except ValueError:
            pass
        packages += ("PyOP2", "firedrake")
        for p in packages:
            clean(p)
            install(p+"/")

        # Ensure pytest is at the latest version
        run_pip(["install", "-U", "pytest"])

with environment(**compiler_env):
    with pipargs("--no-deps"):
        if options["slepc"]:
            build_and_install_slepc()

        if options["opencascade"]:
            build_and_install_pythonocc()

    if args.documentation_dependencies:
        install_documentation_dependencies()

if mode == "update":
    try:
        import firedrake_configuration
        firedrake_configuration.setup_cache_dirs()
        log.info("Clearing just in time compilation caches.")
        from firedrake.tsfc_interface import clear_cache, TSFCKernel
        from pyop2.compilation import clear_cache as pyop2_clear_cache
        print('Removing cached TSFC kernels from %s' % TSFCKernel._cachedir)
        clear_cache()
        pyop2_clear_cache()
    except:                     # noqa: E722
        # Unconditional except in order to avoid upgrade script failures.
        log.error("Failed to clear caches. Try running firedrake-clean.")


os.chdir("../..")

if mode == "install":
    log.info("\n\nSuccessfully installed Firedrake.\n")

    log.info("\nFiredrake has been installed in a python venv. You activate it with:\n")
    log.info("  . %s/bin/activate\n" % firedrake_env)
    log.info("The venv can be deactivated by running:\n")
    log.info("  deactivate\n\n")
    log.info("To upgrade Firedrake activate the venv and run firedrake-update\n")
else:
    log.info("\n\nSuccessfully updated Firedrake.\n")
